package com.sys.ai.service.impl;

import com.agentsflex.core.chain.Chain;
import com.agentsflex.core.chain.ChainEvent;
import com.agentsflex.core.chain.ChainEventListener;
import com.agentsflex.core.chain.event.NodeEndEvent;
import com.alibaba.fastjson2.JSON;
import com.alibaba.fastjson2.JSONArray;
import com.alibaba.fastjson2.JSONObject;
import com.baomidou.mybatisplus.extension.conditions.query.LambdaQueryChainWrapper;
import com.baomidou.mybatisplus.extension.service.impl.ServiceImpl;
import com.sys.ai.domain.AiModel;
import com.sys.ai.domain.AiPlatform;
import com.sys.ai.domain.AiWorkflow;
import com.sys.ai.mapper.AiWorkflowMapper;
import com.sys.ai.service.IAiModelService;
import com.sys.ai.service.IAiPlatformService;
import com.sys.ai.service.IAiWorkflowService;
import com.sys.ai.vo.WorkFlowExecuteVO;
import com.sys.ai.workflow.CustomerKnowledge;
import com.sys.ai.workflow.CustomerLLM;
import com.sys.ai.workflow.CustomerLLMConfig;
import com.sys.common.constant.SSEMessageConstants;
import com.sys.common.sse.SseComponent;
import com.sys.common.utils.SecurityUtils;
import dev.tinyflow.core.Tinyflow;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Lazy;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;
import org.springframework.util.Assert;

import java.util.*;

/**
 * AI流程Service业务层处理
 *
 * @author LGY
 * @date 2025-05-23
 */
@Slf4j
@Service
public class AiWorkflowServiceImpl extends ServiceImpl<AiWorkflowMapper, AiWorkflow> implements IAiWorkflowService {
    @Autowired
    private AiWorkflowMapper aiWorkflowMapper;
    @Lazy
    @Autowired
    private CustomerKnowledge customerKnowledge;
    @Autowired
    private IAiModelService modelService;
    @Autowired
    private IAiPlatformService platformService;
    @Autowired
    private SseComponent sseComponent;

    /**
     * 查询AI流程
     *
     * @param id AI流程主键
     * @return AI流程
     */
    @Override
    public AiWorkflow selectAiWorkflowById(String id) {
        return this.getById(id);
    }

    /**
     * 查询AI流程列表
     *
     * @param aiWorkflow AI流程
     * @return AI流程
     */
    @Override
    public List<AiWorkflow> selectAiWorkflowList(AiWorkflow aiWorkflow) {
        return aiWorkflowMapper.selectAiWorkflowList(aiWorkflow);
    }

    /**
     * 新增AI流程
     *
     * @param aiWorkflow AI流程
     * @return 结果
     */

    @Override
    @Transactional(rollbackFor = RuntimeException.class)
    public void insertAiWorkflow(AiWorkflow aiWorkflow) {
        aiWorkflow.setStatus(0);
        this.save(aiWorkflow);
    }

    /**
     * 修改AI流程
     *
     * @param aiWorkflow AI流程
     * @return 结果
     */
    @Override
    @Transactional(rollbackFor = RuntimeException.class)
    public void updateAiWorkflow(AiWorkflow aiWorkflow) {
        this.updateById(aiWorkflow);
    }

    /**
     * 批量删除AI流程
     *
     * @param ids 需要删除的AI流程主键
     * @return 结果
     */
    @Override
    @Transactional(rollbackFor = RuntimeException.class)
    public void deleteAiWorkflowByIds(String[] ids) {
        this.removeBatchByIds(Arrays.asList(ids));
    }

    /**
     * 执行流程
     *
     * @param graph  流程配置
     * @param params 参数
     * @return
     */
    @Override
    public Map<String, Object> executeWorkFlow(String graph, Map<String, Object> params, String customType) {
        List<WorkFlowExecuteVO> workFlowExecuteVOList = new ArrayList<>();
        JSONObject jsonObject = JSONObject.parseObject(graph);
        Tinyflow tinyflow = new Tinyflow(JSON.toJSONString(jsonObject));
        // 配置 LLM 节点
        JSONArray jsonArray = jsonObject.getJSONArray("nodes");
        Map<String, String> nodeIdType = new HashMap<>();
        jsonArray.forEach(node -> {
            JSONObject nodeObj = (JSONObject) node;
            String nodeType = nodeObj.getString("type");
            nodeIdType.put(nodeObj.getString("id"), nodeType);
            JSONObject dataObj = null;
            switch (nodeType) {
                case "llmNode":
                    // 大模型节点
                    dataObj = nodeObj.getJSONObject("data");
                    String llmId = dataObj.getString("llmId");
                    setLlmProvider(llmId, tinyflow);
                    break;
                case "knowledgeNode":
                    // 知识库节点
                    tinyflow.setKnowledgeProvider(id -> customerKnowledge);
                default:
                    break;
            }

        });

        Chain chain = tinyflow.toChain();
        log.info("流程配置: {}", chain);
        // 添加事件监听
        onFlowListener(chain, workFlowExecuteVOList, nodeIdType, customType);
        // 执行并返回结果
        return chain.executeForResult(params);
    }

    @Override
    public List<AiWorkflow> listAllRelease() {
        return new LambdaQueryChainWrapper<>(baseMapper)
                .eq(AiWorkflow::getStatus, 1)
                .list();
    }

    /**
     * 设置大模型
     *
     * @param llmId    模型表的id
     * @param tinyflow 流程对象
     */
    private void setLlmProvider(String llmId, Tinyflow tinyflow) {
        // 根据llmId 查询对应的模型信息， 【模型编码，模型提供平台，平台授权apiKey】
        AiModel aiModel = modelService.getById(llmId);
        Assert.notNull(aiModel, "模型不存在或已被删除");
        AiPlatform aiPlatform = platformService.getById(aiModel.getPlatformId());
        Assert.notNull(aiPlatform, "LLM提供平台不存在或已被删除");
        CustomerLLMConfig customerLLMConfig = new CustomerLLMConfig(aiModel.getCode(), aiPlatform.getBaseUrl(), aiPlatform.getApiKey());
        CustomerLLM customerLLM = new CustomerLLM(customerLLMConfig);
        tinyflow.setLlmProvider(id -> customerLLM);
    }

    /**
     * 添加事件监听
     *
     * @param chain 链路
     */
    private void onFlowListener(Chain chain, List<WorkFlowExecuteVO> workFlowExecuteVOList, Map<String, String> nodeIdType, String customType) {
        chain.addEventListener(new ChainEventListener() {
            @Override
            public void onEvent(ChainEvent event, Chain chain) {
                log.info(event.toString());
                if (event.getClass().getSimpleName().equals("NodeEndEvent")) {
                    try {
                        NodeEndEvent nodeEndEvent = (NodeEndEvent) event;
                        WorkFlowExecuteVO workFlowExecuteVO = new WorkFlowExecuteVO();
                        workFlowExecuteVO.setNodeId(nodeEndEvent.getNode().getId())
                                .setNodeType(nodeIdType.get(nodeEndEvent.getNode().getId()))
                                .setNodeName(nodeEndEvent.getNode().getName())
                                .setResult(nodeEndEvent.getResult());
                        if (StringUtils.isNotBlank(customType)) {
                            sseComponent.sendMessageAsync(SecurityUtils.getUsername(), SSEMessageConstants.SERVICE_TYPE.WORKFLOW, customType,
                                    JSON.toJSONString(workFlowExecuteVO));
                        }
                        workFlowExecuteVOList.add(workFlowExecuteVO);
                        log.info("节点 - {} - 输出: {}", workFlowExecuteVO.getNodeName(), workFlowExecuteVO.getResult());
                    } catch (Exception e) {
                        log.error("获取节点输出异常", e);
                    }
                }
            }
        });
    }
}